#
# Authors: Simon Vandenhende
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)

import argparse
import cv2
import os
import numpy as np
import sys
import torch
from tensorboardX import SummaryWriter
from utils.config import create_config
from utils.common_config import get_train_dataset, get_transformations,\
                                get_val_dataset, get_train_dataloader, get_val_dataloader,\
                                get_optimizer, get_model, adjust_learning_rate,\
                                get_criterion, get_scheduler
from utils.logger import Logger
from train.train_utils import train_vanilla
from evaluation.evaluate_utils import eval_result
from termcolor import colored
import random

# Parser
parser = argparse.ArgumentParser(description='Vanilla Training')
parser.add_argument('--config_env', default='./configs/env.yml',
                    help='Config file for the environment')
parser.add_argument('--config_exp', default='./configs/vaihingen/hrnet18/height+seg+edge.yml',
                    help='Config file for the experiment')
parser.add_argument('--input_width', default=640)
parser.add_argument('--input_height', default=640)
parser.add_argument('--ssimse', default=False, action='store_true')
parser.add_argument('--silog', default=False, action='store_true')


args = parser.parse_args()

Vaihingen_CATEGORY_NAMES = ['ImSurf', 'Building', 'LowVeg',
                            'Tree', 'Car', 'Clutter']
def main():
    # Retrieve config file
    cv2.setNumThreads(0)
    print('use seed 0')
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    torch.cuda.manual_seed_all(0)
    np.random.seed(0)
    random.seed(0)
    os.environ['PYTHONHASHSEED'] = str(0)
    torch.backends.cudnn.deterministic = True

    p = create_config(args)
    writers = {}
    writers['train'] = SummaryWriter(os.path.join(p['output_dir'], 'log', 'train'))
    writers['val'] = SummaryWriter(os.path.join(p['output_dir'], 'log', 'val'))
    sys.stdout = Logger(os.path.join(p['output_dir'], 'log_file.txt'))
    print(colored(p, 'red'))

    # Get model
    print(colored('Retrieve model', 'blue'))
    model = get_model(p)
    model = torch.nn.DataParallel(model)
    model = model.cuda()

    # Get criterion
    print(colored('Get loss', 'blue'))
    criterion = get_criterion(p)
    criterion.cuda()
    print(criterion)

    # CUDNN
    print(colored('Set CuDNN benchmark', 'blue')) 
    torch.backends.cudnn.benchmark = True

    # Optimizer
    print(colored('Retrieve optimizer', 'blue'))
    optimizer = get_optimizer(p, model)
    print(optimizer)

    # Dataset
    print(colored('Retrieve dataset', 'blue'))
    
    # Transforms 
    train_transforms, val_transforms = get_transformations(p)
    train_dataset = get_train_dataset(p, train_transforms)
    val_dataset = get_val_dataset(p, val_transforms)
    train_dataloader = get_train_dataloader(p, train_dataset)
    val_dataloader = get_val_dataloader(p, val_dataset)
    print('Train samples %d - Val samples %d' %(len(train_dataset), len(val_dataset)))
    print('Train transformations:')
    print(train_transforms)
    print('Val transformations:')
    print(val_transforms)
    print(colored('No checkpoint file at {}'.format(p['checkpoint']), 'blue'))
    start_epoch = 0
    best_result = eval_result(p, val_dataloader, model)

    # Scheduler
    print(colored('Retrieve scheduler', 'blue'))
    lr_scheduler = get_scheduler(p, optimizer, len(train_dataloader))
    print(lr_scheduler)
    
    
    # Main loop
    print(colored('Starting main loop', 'blue'))

    for epoch in range(start_epoch, p['epochs']):
        print(colored('Epoch %d/%d' % (epoch+1, p['epochs']), 'green'))
        print(colored('-'*10, 'green'))

        # Adjust lr
        # lr = adjust_learning_rate(p, optimizer, epoch)
        # print('Adjusted learning rate to {:.5f}'.format(lr))

        # Train 
        print('Train ...')
        eval_train = train_vanilla(p, train_dataloader, model, criterion, optimizer, epoch, writers['train'], lr_scheduler)
        for task_name in eval_train:
            for metric in eval_train[task_name]:
                writers['train'].add_scalar(task_name+'/'+metric, eval_train[task_name][metric], epoch)

        # Evaluate
        # Check if need to perform eval first
        if 'eval_final_10_epochs_only' in p.keys() and p['eval_final_10_epochs_only']: # To speed up -> Avoid eval every epoch, and only test during final 10 epochs.
            if epoch + 1 > p['epochs'] - 10:
                eval_bool = True
            else:
                eval_bool = False
        else:
            eval_bool = True

        # Perform evaluation
        if eval_bool:
            print(colored('#'*18+'   Evaluate ...', 'yellow'))
            curr_result = eval_result(p, val_dataloader, model)
            for task_name in curr_result:
                for metric in curr_result[task_name]:
                    writers['val'].add_scalar(task_name + '/' + metric, curr_result[task_name][metric], epoch)
            # improves, best_result = validate_results(p, curr_result, best_result)
            # if improves:
            #     print('Save new best model')
            #     torch.save(model.state_dict(), p['best_model'])
            print(colored('#'*18+'   Evaluate Done', 'yellow'))

        # Checkpoint
        print('Checkpoint ...')
        torch.save({'optimizer': optimizer.state_dict(), 'model': model.state_dict(), 
                    'epoch': epoch + 1, 'best_result': best_result}, p['checkpoint'])

    # Evaluate best model at the end
    print(colored('Evaluating best model at the end', 'blue'))
    model.load_state_dict(torch.load(p['checkpoint'])['model'])
    eval_stats = eval_result(p, val_dataloader, model)
    writers['train'].close()
    writers['val'].close()


if __name__ == "__main__":
    main()
